#include <iostream>
#include <string>
#include <vector>
#include <stdexcept>
#include <cmath>
#include <array>
#include "sqlite3.h"
#include <opencv2/opencv.hpp>
#include <filesystem> 

namespace fs = std::filesystem; // 文件系统命名空间别名
// 数据库操作封装类
class ChessDatabase {
public:
    ChessDatabase(const std::string& db_path) {
        if (sqlite3_open(db_path.c_str(), &db) != SQLITE_OK) {
            throw std::runtime_error("无法打开数据库: " + std::string(sqlite3_errmsg(db)));
        }
        create_table();
    }

    ~ChessDatabase() {
        sqlite3_close(db);
    }

    struct ChessPiece {
        int num;  // 主键，由数据库自动管理
        double x;
        double y;
        std::string color;
        std::string type;
        std::array<double, 7> huMoments;
    };

    // 添加棋子到数据库
    void addChessPiece(const ChessPiece& piece) {
        const char* sql = "INSERT INTO chess_pieces (x, y, color, type, hu1, hu2, hu3, hu4, hu5, hu6, hu7) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);";

        sqlite3_stmt* stmt;
        if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) != SQLITE_OK) {
            throw std::runtime_error("准备SQL语句失败");
        }

        // 绑定参数（省略 num，因为它是自动递增的）
        sqlite3_bind_double(stmt, 1, piece.x);
        sqlite3_bind_double(stmt, 2, piece.y);
        sqlite3_bind_text(stmt, 3, piece.color.c_str(), -1, SQLITE_STATIC);
        sqlite3_bind_text(stmt, 4, piece.type.c_str(), -1, SQLITE_STATIC);

        for (int i = 0; i < 7; ++i) {
            sqlite3_bind_double(stmt, 5 + i, piece.huMoments[i]);
        }

        if (sqlite3_step(stmt) != SQLITE_DONE) {
            sqlite3_finalize(stmt);
            throw std::runtime_error("插入数据失败: " + std::string(sqlite3_errmsg(db)));
        }

        sqlite3_finalize(stmt);
    }

    // 查询并打印数据库中的表内容
    void printAllChessPieces() {
        const char* sql = "SELECT num, x, y, color, type, hu1, hu2, hu3, hu4, hu5, hu6, hu7 FROM chess_pieces;";
        sqlite3_stmt* stmt;

        if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) != SQLITE_OK) {
            throw std::runtime_error("准备查询语句失败");
        }

        while (sqlite3_step(stmt) == SQLITE_ROW) {
            int num = sqlite3_column_int(stmt, 0);
            double x = sqlite3_column_double(stmt, 1);
            double y = sqlite3_column_double(stmt, 2);
            const char* color = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 3));
            const char* type = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 4));

            std::array<double, 7> huMoments;
            for (int i = 0; i < 7; ++i) {
                huMoments[i] = sqlite3_column_double(stmt, 5 + i);
            }

            // 打印结果
            std::cout << "num: " << num << ", x: " << x << ", y: " << y
                      << ", color: " << color << ", type: " << type << std::endl;
            std::cout << "Hu Moments: ";
            for (double hu : huMoments) {
                std::cout << hu << " ";
            }
            std::cout << std::endl;
        }

        sqlite3_finalize(stmt);
    }
	
std::string matchChessPiece(const ChessPiece& piece) {	
    const char* sql = "SELECT type, hu1, hu2, hu3, hu4, hu5, hu6, hu7 FROM chess_pieces WHERE color = ?;";
    sqlite3_stmt* stmt;
 
    if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) != SQLITE_OK) {
        throw std::runtime_error("准备查询语句失败");
    }
 
    // 绑定颜色参数
    sqlite3_bind_text(stmt, 1, piece.color.c_str(), -1, SQLITE_STATIC);
 
    std::string bestType;
    double minDistance = std::numeric_limits<double>::max();
    int matchCount = 0;
 
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const char* type = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 0));
        std::array<double, 7> dbHuMoments;
        
        // 提取数据库中的Hu矩
        for (int i = 0; i < 7; ++i) {
            dbHuMoments[i] = sqlite3_column_double(stmt, 1 + i);
        }
 
        // 计算欧氏距离
        double distance = 0.0;
        for (int i = 0; i < 7; ++i) {
            double diff = piece.huMoments[i] - dbHuMoments[i];
            distance += diff * diff;
        }
        distance = sqrt(distance);
 
        // 更新最佳匹配
        if (distance < minDistance) {
            minDistance = distance;
            bestType = type;
        }
        matchCount++;
    }
 
    sqlite3_finalize(stmt);
 
    if (matchCount == 0) {
        throw std::runtime_error("未找到同颜色棋子记录");
    }
 
    return bestType;
}
	

private:
    sqlite3* db;

    void create_table() {
        const char* sql = R"(
            CREATE TABLE IF NOT EXISTS chess_pieces (
                num INTEGER PRIMARY KEY AUTOINCREMENT,  -- 自动递增的主键
                x REAL NOT NULL,
                y REAL NOT NULL,
                color TEXT CHECK(color IN ('red', 'black')),
                type TEXT NOT NULL,
                hu1 REAL NOT NULL,
                hu2 REAL NOT NULL,
                hu3 REAL NOT NULL,
                hu4 REAL NOT NULL,
                hu5 REAL NOT NULL,
                hu6 REAL NOT NULL,
                hu7 REAL NOT NULL
            );
        )";

        char* errMsg = nullptr;
        if (sqlite3_exec(db, sql, nullptr, nullptr, &errMsg) != SQLITE_OK) {
            std::string error = "创建表失败: " + std::string(errMsg);
            sqlite3_free(errMsg);
            throw std::runtime_error(error);
        }
    }
};


// 新增的Hu矩提取函数
std::vector<double> extractHuMoments(const cv::Mat& image) {
    // 转换为灰度图像
    cv::Mat grayImg;
    cv::cvtColor(image, grayImg, cv::COLOR_BGR2GRAY);
 
    // 计算Hu矩
    cv::Moments moments = cv::moments(grayImg);
    cv::Mat hu;
    cv::HuMoments(moments, hu);
 
    // 处理Hu矩数值
    std::vector<double> processedHuMoments;
    for (int i = 0; i < 7; ++i) {
        double val = hu.at<double>(i);
        processedHuMoments.push_back(
            std::copysign(1.0, val) * std::log10(std::abs(val) + 1e-10)
        );
    }
 
    return processedHuMoments;
}

// 从图像提取特征的函数（自动识别颜色），以结构体为返回参数的值
ChessDatabase::ChessPiece extractChessPieceFeatures(
    const std::string& imagePath,
    double x,
    double y,
    const std::string& type) // 移除了 color 参数
{
    // 加载图像（彩色模式）
    cv::Mat img = cv::imread(imagePath, cv::IMREAD_COLOR);
    if (img.empty()) {
        throw std::runtime_error("无法加载图像: " + imagePath);
    }
    cv::imshow("Image Display", img);
    cv::waitKey(0);

    // 转换为 HSV 空间
    cv::Mat hsvImg;
    cv::cvtColor(img, hsvImg, cv::COLOR_BGR2HSV);

    // 判断颜色（红色或黑色）
    const int minSaturation = 40;  // 最小饱和度阈值
    const int minValue = 30;       // 最小亮度阈值
    const float redThreshold = 0.3; // 红色像素比例阈值

    int redCount = 0;
    int validPixels = 0;

    // 遍历所有像素
    for (int i = 0; i < hsvImg.rows; ++i) {
        for (int j = 0; j < hsvImg.cols; ++j) {
            cv::Vec3b hsv = hsvImg.at<cv::Vec3b>(i, j);

            // 过滤无效像素（低饱和度或低亮度）
            if (hsv[1] > minSaturation && hsv[2] > minValue) {
               ++validPixels;

                // Hue值范围处理（OpenCV将Hue缩放到0-180）
                if ((hsv[0] <= 10) || (hsv[0] >= 160)) { // 红色范围
                    ++redCount;
                }
            }
        }
    }

    // 打印调试信息
    std::cout << "图像总像素数: " << img.total() << std::endl;
    std::cout << "有效像素数 (validPixels): " << validPixels << std::endl;
    std::cout << "红色像素数 (redCount): " << redCount << std::endl;

/*
    // 有效像素比例不足时抛出异常
    if (validPixels < img.total() * 0.1) {
        throw std::runtime_error("图像有效颜色区域不足");
    }
*/
    // 计算红色像素比例
    float redRatio = static_cast<float>(redCount) / validPixels;
    std::string color = redRatio > redThreshold ? "red" : "black";

    std::cout << "红色像素比例 (redRatio): " << redRatio << std::endl;
    std::cout << "识别颜色: " << color << std::endl;

// 调用封装后的Hu矩提取函数
    std::vector<double> huMoments = extractHuMoments(img);

    // 构造棋子特征对象
    ChessDatabase::ChessPiece piece;
    piece.x = x;
    piece.y = y;
    piece.color = color; // 自动识别的颜色
    piece.type = type;

// 将提取的Hu矩复制到结构体中
    for (int i = 0; i < 7; ++i) {
        piece.huMoments[i] = huMoments[i];
    }
	
    return piece;
}
// 修改后的特征提取函数（增加文件夹遍历逻辑）
void processChessPieces(const std::string& rootDir) {
    try {
        ChessDatabase db("test.db");
        
        // 遍历根目录下的所有子文件夹
        for (const auto& entry : fs::directory_iterator(rootDir)) {
            if (entry.is_directory()) {
                std::string pieceType = entry.path().filename().string(); // 获取子文件夹名称作为类型
                std::cout << "Processing piece type: " << pieceType << std::endl;
 
                // 遍历子文件夹中的所有图片文件
                int fileIndex = 1;
                for (const auto& file : fs::directory_iterator(entry.path())) {
                    if (file.path().extension() == ".png") {
                        std::string imagePath = file.path().string();
                        std::cout << "Processing file: " << imagePath << std::endl;
 
                        try {
                            // 提取特征（自动获取颜色）
                            ChessDatabase::ChessPiece piece = extractChessPieceFeatures(
                                imagePath,
                                1.0,
                                0.0, // y坐标保持为0（可根据需要调整）
                                pieceType
                            );
 
                            // 存储到数据库
                            db.addChessPiece(piece);
                            std::cout << "Stored piece: " << pieceType << " " << fileIndex 
                                      << " successfully!" << std::endl;
                        } catch (const std::exception& e) {
                            std::cerr << "Error processing " << imagePath 
                                      << ": " << e.what() << std::endl;
                        }
 
                        fileIndex++;
                    }
                }
            }
        }
 
        // 打印数据库内容
       // db.printAllChessPieces();
    } catch (const std::exception& e) {
        std::cerr << "Fatal error: " << e.what() << std::endl;
    }
}



// 新增的测试函数
void testChessRecognition(const std::string& imagePath) {
    try {
        // 初始化数据库（假设数据库已存在且有数据）
        ChessDatabase db("test.db");
        
        // 提取测试图像特征（这里可以传入实际坐标，示例用固定值）
        ChessDatabase::ChessPiece testPiece = extractChessPieceFeatures(
            imagePath,
            1.0,  // x坐标示例值
            0.0,  // y坐标示例值
            "unknown" // 初始类型设为unknown
        );
        
        // 执行匹配
        std::string matchedType = db.matchChessPiece(testPiece);
        
        // 输出结果
        std::cout << "\n=== 识别结果 ===" << std::endl;
        std::cout << "图像路径: " << imagePath << std::endl;
        std::cout << "识别类型: " << matchedType << std::endl;
        std::cout << "颜色: " << testPiece.color << std::endl;
        std::cout << "Hu矩特征: ";
        for (double hu : testPiece.huMoments) {
            std::cout << hu << " ";
        }
        std::cout << "\n================================\n" << std::endl;
        
    } catch (const std::exception& e) {
        std::cerr << "识别测试失败: " << e.what() << std::endl;
    }
}

 
int main() {
    // 设置根目录路径（根据你的实际路径修改）
    //std::string rootDirectory = "/mnt/hgfs/aic_jeff/final_project/3.棋局获取显示/3.1棋子检测识别/棋子素材";
    // 开始处理所有棋子
   // processChessPieces(rootDirectory);
   std::string testImagePath = "/mnt/hgfs/aic_jeff/final_project/3.棋局获取显示/3.1棋子检测识别/Project1/Project1/黑炮.png"; 
    testChessRecognition(testImagePath);
   
    return 0;
}